#include <chrono>
#include <atomic>
#include <mutex>
#include <condition_variable>
#include <thread>
#include <functional>
#include <algorithm>
#include <iostream>
#include <vector>
#include <future>

class TokenBucket {
private:
    const double rate_;                         // 令牌生成速率（每秒令牌数）
    const size_t capacity_;                     // 桶的最大容量
    mutable std::atomic<double> tokens_;        // 当前令牌数量（使用原子操作保证线程安全）
    mutable std::atomic<int64_t> last_time_;    // 上次更新令牌数量的时间点
    mutable std::mutex mutex_;                  // 互斥锁，用于保护关键操作
    mutable std::condition_variable cv_;        // 条件变量，用于等待令牌

private:
    /**
     * 获取当前时间戳（微秒）
     * 纳秒转微秒 去除时间信息count
     */ 
    int64_t getCurrentTimeMicroseconds() const {
        return std::chrono::duration_cast<std::chrono::microseconds>(
            std::chrono::steady_clock::now().time_since_epoch()).count();
    }

public:
    /**
     * 构造函数
     * @param rate 令牌生成速率（每秒生成的令牌数）
     * @param capacity 桶的最大容量
     */
    TokenBucket(double rate, size_t capacity) 
        : rate_(rate), capacity_(capacity), tokens_(static_cast<double>(capacity)) {
        last_time_.store(getCurrentTimeMicroseconds());
    }

    /**
     * 更新令牌数量
     */
    void updateTokens() const {
        int64_t now = getCurrentTimeMicroseconds();
        int64_t last = last_time_.load();
        
        // 使用CAS操作更新时间，避免ABA问题
        if (last_time_.compare_exchange_strong(last, now)) {
            // 计算经过的时间（秒）
            double elapsed = (now - last) / 1000000.0;
            // 根据时间增加令牌
            double new_tokens = tokens_.load() + elapsed * rate_;
            // 令牌数量不能超过桶的容量
            tokens_.store(std::min(new_tokens, static_cast<double>(capacity_)));
        }
    }

    /**
     * 尝试获取指定数量的令牌（非阻塞）
     * @param tokens_needed 需要的令牌数量
     * @return 是否获取成功
     */
    bool tryConsume(size_t tokens_needed = 1) {
        updateTokens();
        
        double current_tokens = tokens_.load();
        // compare_exchange_weak可能因为硬件原因出现虚假失败
        while (current_tokens >= tokens_needed) {
            // 即使current_tokens == tokens_，weak版本也可能返回false
            if (tokens_.compare_exchange_weak(current_tokens, current_tokens - tokens_needed)) {
                return true;
            }
            // 虚假失败时，循环会自动重试，无需额外处理
        }
        return false;
    }

    /**
     * 获取指定数量的令牌（阻塞等待）
     * @param tokens_needed 需要的令牌数量
     * @return 是否获取成功
     */
    bool consume(size_t tokens_needed = 1) {
        std::unique_lock<std::mutex> lock(mutex_);
        
        while (true) {
            updateTokens();
            
            double current_tokens = tokens_.load();
            if (current_tokens >= static_cast<double>(tokens_needed)) {
                // 有足够的令牌，直接消费
                if (tokens_.compare_exchange_strong(current_tokens, current_tokens - tokens_needed)) {
                    return true;
                }
            } else {
                // 令牌不足，计算需要等待的时间
                double needed = static_cast<double>(tokens_needed) - current_tokens;
                double wait_time = needed / rate_;
                
                // 等待一段时间后重试
                cv_.wait_for(lock, std::chrono::duration<double>(wait_time), [this, tokens_needed]() {
                    updateTokens();
                    return tokens_.load() >= static_cast<double>(tokens_needed);
                });
            }
        }
    }

    /**
     * 获取当前可用令牌数量
     * @return 可用令牌数
     */
    double availableTokens() const {
        updateTokens();
        return tokens_.load();
    }
};

// 线程安全的输出函数
std::mutex output_mutex;

int main() {
    // 创建一个每秒生成10个令牌，最大容量为50的令牌桶
    TokenBucket bucket(10.0, 50);
    
    std::cout << "Initial tokens: " << bucket.availableTokens() << std::endl;
    
    // 模拟高并发场景
    std::vector<std::future<bool>> futures;
    
    for (int i = 0; i < 20; ++i) {
        futures.push_back(std::async(std::launch::async, [&bucket, i]() {
            // 每个线程尝试获取5个令牌
            bool success = bucket.tryConsume(5);
            
            // 线程安全输出
            {
                std::lock_guard<std::mutex> lock(output_mutex);
                std::cout << "Thread " << i << " " 
                          << (success ? "got tokens" : "failed to get tokens")
                          << " (remaining: " << bucket.availableTokens() << ")" << std::endl;
            }
            
            return success;
        }));
    }
    
    // 统计结果
    int success_count = 0;
    for (auto& future : futures) {
        if (future.get()) {
            success_count++;
        }
    }
    
    std::cout << "\nSummary: " << success_count << "/20 threads got tokens successfully" << std::endl;
    std::cout << "Final tokens: " << bucket.availableTokens() << std::endl;
    
    return 0;
}